from PIL import Image
import sys
import os
import urllib
import tensorflow.contrib.tensorrt as trt
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import tensorflow as tf
import numpy as np
import time
from tf_trt_models.detection import download_detection_model, build_detection_graph

_GPU_MEM_FRACTION = 0.50
MODEL = 'ssd_mobilenet_v1_coco'
DATA_DIR = './data/'
CONFIG_FILE = MODEL + '.config'   # ./data/ssd_inception_v2_coco.config 
CHECKPOINT_FILE = 'model.ckpt'    # ./data/ssd_inception_v2_coco/model.ckpt
IMAGE_PATH = './data/huskies.jpg'

#config_path, checkpoint_path = download_detection_model(MODEL, 'data')
'''
frozen_graph, input_names, output_names = build_detection_graph(
    config='./data/ssd_mobilenet_v1_coco_2018_01_28/pipeline.config',
    checkpoint='./data/ssd_mobilenet_v1_coco_2018_01_28/model.ckpt',
    score_threshold=0.3,
    batch_size=1
)

print(input_names,output_names)

trt_graph = trt.create_inference_graph(
    input_graph_def=frozen_graph,
    outputs=output_names,
    max_batch_size=1,
    max_workspace_size_bytes=1 << 25,
    precision_mode='INT8',
    minimum_segment_size=50
)


with open('./data/trt.pb', 'wb') as f:
    f.write(trt_graph.SerializeToString())
'''
pb_fname = "./data/trt.pb"
detection_graph = tf.Graph()
with detection_graph.as_default():
	od_graph_def = tf.GraphDef()
	with tf.gfile.GFile(pb_fname, 'rb') as fid:
		serialized_graph = fid.read()
		od_graph_def.ParseFromString(serialized_graph)
		tf.import_graph_def(od_graph_def, name='')


tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
tf_sess = tf.Session(graph=detection_graph,config=tf_config)

#tf.import_graph_def(detection_graph, name='')
tf_input = tf_sess.graph.get_tensor_by_name('image_tensor:0')
tf_scores = tf_sess.graph.get_tensor_by_name('detection_scores:0')
tf_boxes = tf_sess.graph.get_tensor_by_name('detection_boxes:0')
tf_classes = tf_sess.graph.get_tensor_by_name('detection_classes:0')
tf_num_detections = tf_sess.graph.get_tensor_by_name('num_detections:0')

image = Image.open(IMAGE_PATH)
image_resized = np.array(image)
image_resized = np.expand_dims(image_resized, axis=0)
#image = np.array(image)
'''
plt.imshow(image)

image_resized = np.array(image.resize((300, 300)))
image = np.array(image)

scores, boxes, classes, num_detections = tf_sess.run([tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={
    tf_input: image_resized[None, ...]
})

boxes = boxes[0] # index by 0 to remove batch dimension
scores = scores[0]
classes = classes[0]
num_detections = int(num_detections[0])

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)

ax.imshow(image)
'''
# plot boxes exceeding score threshold

num_samples = 50

t0 = time.time()
for i in range(num_samples):
    scores, boxes, classes, num_detections = tf_sess.run([tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={
        tf_input: image_resized
    })
t1 = time.time()
print('Average runtime: %f seconds' % (float(t1 - t0)/num_samples))

tf_sess.close()
